package com.cctc.rds.raft.core.rpc.log;

import com.cctc.rds.raft.core.RaftCluster;
import com.cctc.rds.raft.core.log.Log;
import com.cctc.rds.raft.core.log.LogService;
import com.cctc.rds.raft.core.node.Node;
import com.cctc.rds.raft.core.rpc.RpcClientRouter;
import com.cctc.rds.raft.core.util.CountDownWaiter;
import com.cctc.rds.raft.core.util.IdUtils;

import java.io.Closeable;
import java.io.IOException;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

public class LogRpcScheduler implements Closeable {
    private int nodeId;
    private LogService logService;
    private RpcClientRouter<LogRpc> router;
    private ExecutorService logWorkerGroup, heartBeatWorkerGroup;

    public LogRpcScheduler(int nodeId, LogService logService, RpcClientRouter<LogRpc> router) {
        this.nodeId = nodeId;
        this.logService = logService;
        this.router = router;
        this.logWorkerGroup = newExecutorService("log-rpc-worker", RaftCluster.nodeCount() - 1, 1);
        this.heartBeatWorkerGroup = newExecutorService("heartbeat-rpc-worker", RaftCluster.nodeCount() - 1, 1);
    }

    private ExecutorService newExecutorService(String name, int threads, int queueCapacity) {
        return new ThreadPoolExecutor(
                threads, threads,
                60, TimeUnit.SECONDS,
                new LinkedBlockingDeque<>(queueCapacity),
                r -> new Thread(r, name + "-" + IdUtils.newId()),
                new ThreadPoolExecutor.CallerRunsPolicy()
        );
    }

    public LogRpcRequest newLogRpcRequest() {
        Node curNode = RaftCluster.getNode(nodeId);
        LogRpcRequest request = LogRpcRequest.builder()
                .nodeId(curNode.getId())
                .term(curNode.getCurTerm())
                .build();
        Log log = logService.index(curNode.curCommitLogIndex());
        if (log == null) {
            request.setLastCommitLogIndex(-1);
            request.setLastCommitLogTerm(-1);
        } else {
            request.setLastCommitLogIndex(log.getIndex());
            request.setLastCommitLogTerm(log.getTerm());
        }
        return request;
    }

    public boolean append(LogRpcRequest request) {
        Node leader = RaftCluster.getNode(nodeId);
        Set<Node> nodes = RaftCluster.getAllNodes();
        int count = RaftCluster.nodeCount();
        final CountDownWaiter waiter = new CountDownWaiter(count - 1);
        for (Node toNode : nodes) {
            if (toNode.getId().equals(leader.getId())) continue;
            request.setFollowerCommitLogIndex(toNode.curCommitLogIndex());
            logWorkerGroup.execute(() -> {
                try {
                    LogRpcResponse response = router.route(toNode.getIpPort()).append(request);
                    if (response.isSuccess()) {
                        waiter.countDownSuccess();
                    } else if (response.getTerm() <= request.getTerm()) {
                        LogRpcRequest prevRequest = request.clone();
                        long lastCommitLogIndex = request.getLastCommitLogIndex();
                        long prevCommitLogIndex = lastCommitLogIndex - 1;
                        while (prevCommitLogIndex >= -1) {
                            if (prevCommitLogIndex == -1) {
                                prevRequest.setLastCommitLogIndex(-1);
                                prevRequest.setLastCommitLogTerm(-1);
                                prevRequest.setLogs(logService.range(0, lastCommitLogIndex));
                            } else {
                                Log log = logService.index(prevCommitLogIndex);
                                prevRequest.setLastCommitLogIndex(log.getIndex());
                                prevRequest.setLastCommitLogTerm(log.getTerm());
                                prevRequest.setLogs(logService.range(prevCommitLogIndex, lastCommitLogIndex));
                            }
                            LogRpcResponse prevResponse = router.route(toNode.getIpPort()).append(prevRequest);
                            if (prevResponse.isSuccess()) {
                                response = router.route(toNode.getIpPort()).append(request);
                                if (response.isSuccess()) {
                                    waiter.countDownSuccess();
                                }
                                break;
                            }
                            prevCommitLogIndex--;
                        }
                    } else {
                        // leader is old, redirect to heartbeat.
                    }
                } catch (Exception e) {
                    waiter.countDownException();
                } finally {
                    waiter.countDown();
                }
            });
        }
        int half = (count / 2) + 1;
        waiter.await();
        return waiter.getCountSuccess() >= half;
    }

    public boolean heartbeat() {
        Node leader = RaftCluster.getNode(nodeId);
        Set<Node> nodes = RaftCluster.getAllNodes();
        int count = RaftCluster.nodeCount();
        final CountDownWaiter waiter = new CountDownWaiter(count - 1);
        final AtomicBoolean success = new AtomicBoolean(Boolean.TRUE);
        for (Node toNode : nodes) {
            if (toNode.getId().equals(leader.getId())) continue;
            LogRpcRequest request = newLogRpcRequest();
            request.setFollowerCommitLogIndex(toNode.curCommitLogIndex());
            heartBeatWorkerGroup.execute(() -> {
                try {
                    LogRpcResponse response = router.route(toNode.getIpPort()).append(request);
                    if (!response.isSuccess() && response.getTerm() > request.getTerm()) {
                        success.set(Boolean.FALSE);
                    }
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    waiter.countDown();
                }
            });
        }
        waiter.await();
        return success.get();
    }

    public void commit(Log log) {
        Node leader = RaftCluster.getNode(nodeId);
        Set<Node> nodes = RaftCluster.getAllNodes();
        for (Node toNode : nodes) {
            if (toNode.getId().equals(leader.getId())) continue;
            logWorkerGroup.execute(() -> {
                try {
                    router.route(toNode.getIpPort()).apply(log.getIndex(), log.getTerm());
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });
        }
    }

    @Override
    public void close() throws IOException {
        logWorkerGroup.shutdown();
        heartBeatWorkerGroup.shutdownNow();
    }
}
